package demoMod.stsconsole.utils;

import demoMod.stsconsole.annotations.*;
import demoMod.stsconsole.exceptions.MissingParamTypesException;
import demoMod.stsconsole.exceptions.PatchingException;
import demoMod.stsconsole.patches.*;
import demoMod.stsconsole.patches.lib.EnumBusterReflect;
import demoMod.stsconsole.patches.lib.ReflectionHelper;
import javassist.*;

import java.io.File;
import java.io.IOException;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.URL;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class ModLoader {
    private static final String CLASS_FILE_SUFFIX = ".class";
    private static final ClassPool cp = ClassPool.getDefault();
    private static final TreeSet<PatchInfo> patchInfos = new TreeSet<>(new PatchInfoComparator());
    private static final Map<Class<?>, EnumBusterReflect> enumBusterMap = new HashMap<>();
    public static final List<CtClass> ctClasses = new ArrayList<>();

    private final URL[] urls;
    private static final Set<String> loadedJarFiles = new HashSet<>();

    public ModLoader(String modFolder) {
        File file = new File(modFolder);
        if (!file.exists()) {
            file.mkdirs();
        }
        File[] jarFiles = file.listFiles((dir, name) -> name.endsWith(".jar"));
        if (jarFiles == null) {
            jarFiles = new File[0];
        }
        urls = new URL[jarFiles.length];
        int index = 0;
        for (File jarFile : jarFiles) {
            try {
                if (!loadedJarFiles.contains(jarFile.getAbsolutePath())) {
                    urls[index] = jarFile.toURI().toURL();

                    cp.appendClassPath(jarFile.getAbsolutePath());
                    loadedJarFiles.add(jarFile.getAbsolutePath());
                }
            } catch (IOException | NotFoundException e) {
                e.printStackTrace();
            }
            index++;
        }
    }

    public void loadMods() {
        try {
            for (URL url : urls) {
                scanJ(url);
                PatchInfo.nextMod();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 扫描Jar包下所有class
     * @param url jar-url路径
     * @throws IOException 当文件读写发生错误时抛出此异常
     * @throws ClassNotFoundException 如果加载类时没有找到指定的类，就会抛出此异常
     */
    private void scanJ(URL url) throws Exception {
        if (url == null) return;
        JarFile jarFile = new JarFile(url.getFile());
        List<String> classNames = new ArrayList<>();
        // 遍历Jar包
        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry jarEntry = entries.nextElement();
            String fileName = jarEntry.getName();

            if (jarEntry.isDirectory()) {
                continue;
            }

            if (fileName.endsWith(CLASS_FILE_SUFFIX)) {
                String className = fileName.substring(0, fileName.indexOf('.')).replace('/', '.');
                classNames.add(className);
            }
        }
        injectPatches(classNames);
    }

    /**
     * 将patch注入到指定的类中
     * @param classNames 类名
     */
    private void injectPatches(List<String> classNames) throws Exception {
        for (String className : classNames) {
            CtClass ctClass = cp.get(className);
            for (CtField ctField : ctClass.getDeclaredFields()) {
                if (ctField.hasAnnotation(SpireEnum.class)) {
                    Class cls = ctClass.toClass();
                    SpireEnum spireEnum = (SpireEnum) ctField.getAnnotation(SpireEnum.class);
                    Field field = cls.getDeclaredField(ctField.getName());
                    String enumName = field.getName();
                    if (!spireEnum.name().isEmpty()) {
                        enumName = spireEnum.name();
                    }

                    EnumBusterReflect buster;
                    if (enumBusterMap.containsKey(field.getType())) {
                        buster = enumBusterMap.get(field.getType());
                    } else {
                        buster = new EnumBusterReflect(ModLoader.class.getClassLoader(), field.getType());
                        enumBusterMap.put(field.getType(), buster);
                    }
                    Enum<?> enumValue = buster.make(enumName);
                    buster.addByValue(enumValue);
                    try {
                        Field constantField = field.getType().getField(enumName);
                        ReflectionHelper.setStaticFinalField(constantField, enumValue);
                    } catch (NoSuchFieldException ignored) {
                    }

                    field.setAccessible(true);
                    field.set(null, enumValue);
                }
            }
            if (ctClass.hasAnnotation(SpirePatch.class)) {
                List<SpirePatch> patchArr = new ArrayList<>();
                SpirePatches patches = (SpirePatches) ctClass.getAnnotation(SpirePatches.class);
                if (patches != null) {
                    Collections.addAll(patchArr, patches.value());
                } else {
                    SpirePatch patch = (SpirePatch) ctClass.getAnnotation(SpirePatch.class);
                    if (patch != null) {
                        patchArr.add(patch);
                    }
                }

                for (SpirePatch patch : patchArr) {
                    CtClass ctClsToPatch = null;
                    try {
                        if (!patch.cls().isEmpty()) {
                            ctClsToPatch = cp.get(patch.cls());
                        }
                    } catch (NotFoundException e) {
                        if (patch.optional()) {
                            continue;
                        }
                        throw new PatchingException(ctClass.getName(), e);
                    }
                    if (ctClsToPatch == null) {
                        throw new PatchingException(ctClass, "No class defined to patch. Must define cls in @SpirePatch.");
                    }
                    if (!ctClasses.contains(ctClsToPatch)) {
                        ctClasses.add(ctClsToPatch);
                    }
                    CtBehavior ctMethodToPatch = null;
                    try {
                        CtClass[] ctParamTypes = patchParamTypes(patch);
                        switch (patch.method()) {
                            case SpirePatch.CONSTRUCTOR:
                                if (ctParamTypes == null) {
                                    CtConstructor[] constructors = ctClsToPatch.getDeclaredConstructors();
                                    if (constructors.length == 1) {
                                        ctMethodToPatch = constructors[0];
                                    } else {
                                        throw new MissingParamTypesException(ctClass, patch);
                                    }
                                } else {
                                    ctMethodToPatch = ctClsToPatch.getDeclaredConstructor(ctParamTypes);
                                }
                                break;
                            case SpirePatch.CL_INIT:
                                ctMethodToPatch = ctClsToPatch.getClassInitializer();
                                if (ctMethodToPatch == null) {
                                    System.out.println("No class initializer, making one");
                                    ctMethodToPatch = ctClsToPatch.makeClassInitializer();
                                }
                                break;
                            case SpirePatch.CLASS:
                                patchInfos.add(new ClassPatchInfo(ctClsToPatch, ctClass));
                                break;
                            default:
                                if (ctParamTypes == null) {
                                    CtMethod[] methods = ctClsToPatch.getDeclaredMethods(patch.method());
                                    if (methods.length == 1) {
                                        ctMethodToPatch = methods[0];
                                    } else if (methods.length == 0) {
                                        throw new NoSuchMethodException(String.format("Patch %s:\nNo method named [%s] found on\nclass [%s]",
                                                ctClass.getName(),
                                                patch.method(),
                                                patchClassName(patch)
                                        ));
                                    } else {
                                        throw new MissingParamTypesException(ctClass, patch);
                                    }
                                } else {
                                    ctMethodToPatch = ctClsToPatch.getDeclaredMethod(patch.method(), ctParamTypes);
                                }
                                break;
                        }
                    } catch (NotFoundException e) {
                        throw new NoSuchMethodException(String.format("Patch %s:\nNo method [%s(%s)] found on\nclass [%s]",
                                ctClass.getName(),
                                patch.method(),
                                patchParamTypesString(patch),
                                patchClassName(patch)
                        ));
                    } catch (MissingParamTypesException e) {
                        e.printStackTrace();
                    }
                    if (ctMethodToPatch == null) continue;
                    for (CtMethod m : ctClass.getDeclaredMethods()) {
                        PatchInfo p = null;
                        if (m.getName().equals("Prefix")) {
                            p = new PrefixPatchInfo(ctMethodToPatch, m).setSpirePatch(patch);
                        } else if (m.getName().equals("Postfix")) {
                            p = new PostfixPatchInfo(ctMethodToPatch, m).setSpirePatch(patch);
                        } else if (m.getName().equals("Locator")) {
                            continue;
                        } else if (m.hasAnnotation(SpireInsertPatch.class)) {
                            SpireInsertPatch insertPatch = (SpireInsertPatch) m.getAnnotation(SpireInsertPatch.class);

                            LocatorInfo locatorInfo = null;
                            if (insertPatch != null && !insertPatch.locator().equals(SpireInsertPatch.NONE.class)) {
                                locatorInfo = new LocatorInfo(ctMethodToPatch, cp.get(insertPatch.locator().getName()).toClass());
                            }

                            if (!isInsertPatchValid(insertPatch, locatorInfo)) {
                                throw new PatchingException(m, "SpireInsertPatch missing line number! Must specify either loc, rloc, locs, rlocs, or a Locator");
                            }

                            List<InsertPatchInfo.LineNumberAndPatchType> locs = new ArrayList<>();

                            if (locatorInfo != null) {
                                int[] abs_locs = locatorInfo.findLines();
                                if (abs_locs.length < 1) {
                                    throw new PatchingException(m, "Locator must locate at least 1 line!");
                                }
                                for (int abs_loc : abs_locs) {
                                    locs.add(new InsertPatchInfo.LineNumberAndPatchType(abs_loc));
                                }
                            }

                            if (insertPatch.loc() >= 0) {
                                locs.add(new InsertPatchInfo.LineNumberAndPatchType(insertPatch.loc()));
                            }
                            if (insertPatch.rloc() >= 0) {
                                locs.add(new InsertPatchInfo.LineNumberAndPatchType(
                                        ctMethodToPatch.getMethodInfo().getLineNumber(0) + insertPatch.rloc(), insertPatch.rloc()));
                            }
                            for (int i = 0; i < insertPatch.locs().length; i++) {
                                locs.add(new InsertPatchInfo.LineNumberAndPatchType(insertPatch.locs()[i]));
                            }
                            for (int i = 0; i < insertPatch.rlocs().length; i++) {
                                locs.add(new InsertPatchInfo.LineNumberAndPatchType(
                                        ctMethodToPatch.getMethodInfo().getLineNumber(0) + insertPatch.rlocs()[i], insertPatch.rlocs()[i]));
                            }

                            p = new InsertPatchInfo(insertPatch, locs, ctMethodToPatch, m).setSpirePatch(patch);
                        } else if (m.getName().equals("Instrument")) {
                            p = new InstrumentPatchInfo(ctMethodToPatch, findInstrumentMethod(cp.get(className).toClass(), m.getName())).setSpirePatch(patch);
                        } else if (m.getName().equals("Raw")) {
                            p = new RawPatchInfo(ctMethodToPatch, findRawMethod(cp.get(className).toClass(), m.getName())).setSpirePatch(patch);
                        }

                        if (p != null) {
                            if (!Modifier.isPublic(m.getModifiers())) {
                                m.setModifiers(Modifier.setPublic(m.getModifiers()));
                            }
                            patchInfos.add(p);
                        }
                    }
                }
            }
            if (!ctClasses.contains(ctClass)) {
                ctClasses.add(ctClass);
            }
        }
        System.out.println("Injecting patches...");

        for (PatchInfo p : patchInfos) {
            p.debugPrint();

            try {
                p.doPatch();
            } catch (Exception e) {
                System.out.println();
                p.debugPrint();
                throw e;
            }

            System.out.println();
        }
        patchInfos.clear();
        System.out.println("Done.");
    }

    /**
     * 加载所有mod类
     */
    public static void loadClass() throws CannotCompileException, NoSuchMethodException, InvocationTargetException, IllegalAccessException {
        List<Class> invokeLater = new ArrayList<>();
        for (CtClass ctClass : ctClasses) {
            try {
                Class cls = ctClass.toClass(ModLoader.class.getClassLoader(), null);
                if (cls.isAnnotationPresent(SpireInitializer.class)) {
                    invokeLater.add(cls);
                }
                System.out.println(cls.getName());
            } catch (CannotCompileException e) {
                if (e.getCause() instanceof ClassFormatError) {
                    System.out.println("load " + ctClass.getName() + " failed cause duplicate class definition.");
                } else {
                    throw e;
                }
            }
            ctClass.detach();
        }
        ctClasses.clear();
        for (Class cls : invokeLater) {
            Method method = cls.getDeclaredMethod("initialize");
            method.setAccessible(true);
            method.invoke(null);
        }
    }

    private static String patchParamTypesString(SpirePatch patch) {
        if (patch.paramtypes().length == 1 && "void".equals(patch.paramtypes()[0])) {
            String[] def = {"DEFAULT"};
            if (Arrays.equals(patch.paramtypes(), def))
                return "";

            return String.join(", ", patch.paramtypes());
        } else {
            String[] tmp = new String[patch.paramtypes().length];
            System.arraycopy(patch.paramtypes(), 0, tmp, 0, patch.paramtypes().length);
            return String.join(", ", tmp);
        }
    }

    private static String patchClassName(SpirePatch patch) {
        return patch.cls();
    }

    private static CtClass[] patchParamTypes(SpirePatch patch) throws NotFoundException {
        String[] def = {"DEFAULT"};
        if (Arrays.equals(patch.paramtypes(), def)) {
            return null;
        }

        return cp.get(patch.paramtypes());
    }

    private static boolean isInsertPatchValid(SpireInsertPatch insertPatch, LocatorInfo locatorInfo) {
        if (locatorInfo != null) {
            return true;
        }
        if (insertPatch != null) {
            return insertPatch.loc() != -1 || insertPatch.rloc() != -1
                    || insertPatch.locs().length != 0 || insertPatch.rlocs().length != 0;
        }
        return false;
    }

    private static Method findInstrumentMethod(Class<?> cls, String name) throws NoSuchMethodException {
        for (Method m : cls.getDeclaredMethods()) {
            if (m.getName().equals(name) && m.getParameterCount() == 0) {
                m.setAccessible(true);
                return m;
            }
        }
        throw new NoSuchMethodException();
    }

    private static Method findRawMethod(Class<?> cls, String name) throws NoSuchMethodException {
        for (Method m : cls.getDeclaredMethods()) {
            if (m.getName().equals(name) && m.getParameterCount() == 1) {
                m.setAccessible(true);
                return m;
            }
        }
        throw new NoSuchMethodException();
    }
}
